#####
# Replication for: "Can Political Speech Foster Tolerance of Immigrants?" by Schleiter, Tavits, and Ward.
# Figure S.3
#####

library(here)
library(data.table)
library(ggplot2)

# source the custom functions
source("functions.R")

# load the pooled data if not already in workspace
if(!exists("pooled")){
  pooled <- fread(here("data", "pooled.csv"))
}

# fit models --------------------------------------------------------------

#1. gender
bal_gend <- noCov(female, sample = "Pooled")

#2. age
pooled[ , age_center_pool := (age - mean(age, na.rm = T))/sd(age, na.rm = T)]
bal_age_center <- noCov(age_center_pool, sample = "Pooled")

#3. race
bal_white <- noCov(I(race2 == "White/Caucasian"), sample = "Pooled")
bal_asian <- noCov(I(race2 == "Asian"), sample = "Pooled")
bal_afam <- noCov(I(race2 == "Black/African American"), sample = "Pooled")
bal_hisp <- noCov(I(race2 == "Hispanic/Latino"), sample = "Pooled")
bal_other_race <- noCov(I(race2 == "Other (please specify)"), sample = "Pooled")

#4. Education
bal_lessthan <- noCov(I(edu2 == "Less than high school"), sample = "Pooled")
bal_highSchool <- noCov(I(edu2 == "High school graduate/ GED"), sample = "Pooled")
bal_2year <- noCov(I(edu2 == "2 year degree"), sample = "Pooled")
bal_someCollege <- noCov(I(edu2 == "Some college"), sample = "Pooled")
bal_4year <- noCov(I(edu2 == "4 year degree"), sample = "Pooled")
bal_masters <- noCov(I(edu2 == "Masters degree"), sample = "Pooled")
bal_prof <- noCov(I(edu2 == "Professional Degree (JD, MD)"), sample = "Pooled")
bal_doc <- noCov(I(edu2 == "Doctorate"), sample = "Pooled")

#5. Party identification
bal_dem <- noCov(I(pid2 == "Democrat"), sample = "Pooled")
bal_rep <- noCov(I(pid2 == "Republican"), sample = "Pooled")
bal_ind <- noCov(I(pid2 == "Independent"), sample = "Pooled")
bal_other_pid <- noCov(I(pid2 == "Something else"), sample = "Pooled")

#6. News attention
bal_new <- noCov(news2, sample = "Pooled")

#7. Region
bal_NE <- noCov(I(region2 == "Northeast"), sample = "Pooled")
bal_MW <- noCov(I(region2 == "Midwest"), sample = "Pooled")
bal_S <- noCov(I(region2 == "South"), sample = "Pooled")
bal_W <- noCov(I(region2 == "West"), sample = "Pooled")

#8. Wave
bal_w1 <- noCov(I(wave == "1"), sample = "Pooled")
bal_w2 <- noCov(I(wave == "2"), sample = "Pooled")
bal_w3 <- noCov(I(wave == "3"), sample = "Pooled")
bal_replication <- noCov(I(wave == "Replication"), sample = "Pooled")

# Plotting prep -----------------------------------------------------------

bal_models_Pool <- 
  list(
    bal_age_center,
    bal_gend,
    bal_white, bal_asian, bal_afam, bal_hisp, bal_other_race,
    bal_lessthan, bal_highSchool, bal_2year, bal_someCollege, bal_4year, bal_masters, bal_prof, bal_doc,
    bal_dem, bal_rep, bal_ind, bal_other_pid,
    bal_new,
    bal_NE, bal_MW, bal_S, bal_W,
    bal_w1, bal_w2, bal_w3, bal_replication
  )

bal_plot_Pool <- do.call("rbind", lapply(bal_models_Pool, function(mod){
  data.frame(extractPlotInfo(mod), 1, c("Common Humanity", "Countering Stereotypes","Norms"), stringsAsFactors = F)
}))

colnames(bal_plot_Pool) <- c("est", "lower95", "upper95", "lower90", "upper90", "x", "D")

bp_labs <-
  c("Age (centered to\nmean=0, SD=1):",
    "",
    "Female:",
    "",
    "Race:",
    "   White",
    "   Asian",
    "   African American/Black",
    "   Hispanic/Latino",
    "   Other",
    "",
    "Education:",
    "   Less than high school",
    "   High school grad./GED",
    "   2-year degree",
    "   Some college",
    "   4-year degree",
    "   Masters degree",
    "   Professional Degree",
    "   Doctorate",
    "",
    "Party Identification:",
    "   Democrat",
    "   Republican",
    "   Independent",
    "   Something else",
    "",
    "News Attention (1-7):",
    "",
    "Region:",
    "   Northeast",
    "   Midwest",
    "   South",
    "   West"
  ) 

bp_labs_Pool <- c(bp_labs, "", "Study:", "   Main, W1", "   Main, W2", "   Main, W3", "   Replication")

bal_plot_Pool$x <- rev(rep(c(1:4, c(6:9, 12, 14:17, 20:27, 30:34, 37, 39) + 1), each = 3))
bal_plot_Pool$D <- factor(bal_plot_Pool$D, levels = c("Norms", "Countering Stereotypes", "Common Humanity"))

# Plot and save -----------------------------------------------------------

fig_s3 <- ggplot(bal_plot_Pool, aes(x = est, y = x, xmin = lower95, xmax = upper95, group = D, color = D, shape = D)) +
  geom_vline(xintercept = 0, linetype = "dashed") +
  ggstance::geom_pointrangeh(fill = "white") +
  scale_y_continuous(breaks = length(bp_labs_Pool):1, labels = bp_labs_Pool, limits = c(.5, length(bp_labs_Pool) + 1), expand = expansion(0,0)) +
  ylab(NULL) +
  xlab("Regression Coefficients from Balance Tests") +
  scale_color_grey(name = "Treatment:", guide = guide_legend(reverse = T), start = .6, end = 0.2) +
  scale_shape_manual(values = c(17, 21, 15), name = "Treatment:", guide = guide_legend(reverse = T)) +
  theme(axis.text.y = element_text(hjust = 0, face = ifelse(length(bp_labs_Pool):1 %in% unique(bal_plot_Pool$x)[-c(1,2,20)] , "plain", "bold")),
        legend.position = "none",
        panel.grid.minor = element_blank()) +
  facet_wrap( ~ D, ncol = 3)

ggsave(here("SI", "figure_s3.pdf"), fig_s3, height = 8, width = 8)